# -*- coding: utf-8 -*-
import asyncio
import json
import time
import logging
from typing import Dict, Any, Set, Optional
import websockets
from websockets.asyncio.server import ServerConnection
from websockets.exceptions import ConnectionClosed
import struct

from common.landmarks import NAME_TO_IDX, landmark_to_array
from recognition.manager import ActionRecognitionManagerRef
from common.config import HOST, PORT

CLIENTS: Set[ServerConnection] = set()
SUBSCRIBERS: Set[ServerConnection] = set()
_manager_ref = ActionRecognitionManagerRef()

logger = logging.getLogger(__name__)

# 二进制视频帧头（与前端一致）
# | magic(2B='RB') | ver(1B=1) | width(2B) | height(2B) | channels(1B=4) | fmt(1B=0=RGBA8) |
HEADER_FMT = "<2sBHHBB"
MAGIC = b"RB"
VERSION = 1
FMT_RGBA8 = 0
HEADER_LEN = struct.calcsize(HEADER_FMT)

def attach_manager(manager):
    _manager_ref.set(manager)

def get_clients():
    return CLIENTS

def _is_ws_closed(ws) -> bool:
    """
    - 优先用 ws.state（新版本）
    - 回退用 ws.closed（旧版本）或 ws.close_code
    """
    # 新版：ServerConnection 有 state（ConnectionState）
    state = getattr(ws, "state", None)
    if state is not None:
        # ConnectionState.OPEN/CLOSING/CLOSED；用 name 或字符串比较皆可
        name = getattr(state, "name", None)
        if isinstance(name, str):
            return name in ("CLOSING", "CLOSED")
        # 保险起见，字符串化判断
        s = str(state).upper()
        return ("CLOSING" in s) or ("CLOSED" in s)

    # 旧版：WebSocketServerProtocol 有 closed(bool)
    closed_attr = getattr(ws, "closed", None)
    if isinstance(closed_attr, bool):
        if closed_attr:
            return True

    # 兜底：如果有 close_code 也视作已关闭
    return getattr(ws, "close_code", None) is not None

async def safe_send(ws, data_str: str) -> bool:
    """
      - 发送前判断已关闭则直接移除
      - 捕获 ConnectionClosed / CancelledError
      - 其他错误也移除（防止脏连接）
    """
    if _is_ws_closed(ws):
        CLIENTS.discard(ws)
        return False
    try:
        await ws.send(data_str)
        return True
    except (ConnectionClosed, asyncio.CancelledError) as e:
        CLIENTS.discard(ws)
        logger.info(
            "safe_send: connection closed during send: code=%s reason=%s",
            getattr(e, "code", None),
            getattr(e, "reason", "")
        )
        return False
    except Exception:
        CLIENTS.discard(ws)
        logger.exception("safe_send: unexpected send error, removed client")
        return False
        
async def broadcast_json(message: Dict[str, Any]):
    if not CLIENTS:
        return
    data = json.dumps(message, ensure_ascii=False)
    for ws in list(CLIENTS):
        await safe_send(ws, data)

async def safe_send_binary(ws: ServerConnection, data: bytes) -> bool:
    if _is_ws_closed(ws):
        CLIENTS.discard(ws)
        SUBSCRIBERS.discard(ws)
        return False
    try:
        await ws.send(data)  # websockets 库会根据 bytes 自动按二进制发
        return True
    except (ConnectionClosed, asyncio.CancelledError) as e:
        CLIENTS.discard(ws)
        SUBSCRIBERS.discard(ws)
        logger.info(
            "safe_send_binary: connection closed during send: code=%s reason=%s",
            getattr(e, "code", None),
            getattr(e, "reason", "")
        )
        return False
    except Exception:
        CLIENTS.discard(ws)
        SUBSCRIBERS.discard(ws)
        logger.exception("safe_send_binary: unexpected send error, removed client")
        return False

async def send_video_frame_to_subscribers(width: int, height: int, channels: int, fmt_code: int, payload: bytes):
    """
    将一帧视频（已是RGB8字节阵列）发送给所有订阅者。
    - width/height/channels/fmt_code 与 HEADER_FMT 对齐
    """
    if not SUBSCRIBERS:
        return
    header = struct.pack(HEADER_FMT, MAGIC, VERSION, width, height, channels, fmt_code)
    packet = header + payload
    # 顺序发送，避免在当前事件循环里引入复杂并行
    subs = list(SUBSCRIBERS)
    for ws in subs:
        await safe_send_binary(ws, packet)

# 默认 9 点集合（与后端索引映射保持一致）
DEFAULT_9 = [
    "NOSE",
    "LEFT_SHOULDER", "RIGHT_SHOULDER",
    "LEFT_WRIST", "RIGHT_WRIST",
    "LEFT_HIP", "RIGHT_HIP",
    "LEFT_ANKLE", "RIGHT_ANKLE",
]

def collect_joints_from_landmarks(lms, joints):
    """
    从 landmarks 中按 joints（名称或索引）收集坐标，输出 {key: [x,y,z,vis]}。
    - 名称 key: 原样返回
    - 索引 key: 转为字符串
    """
    out = {}
    for j in joints:
        if isinstance(j, int):
            idx = j
            key = str(j)
        else:
            if j not in NAME_TO_IDX:
                continue
            idx = NAME_TO_IDX[j]
            key = j
        if 0 <= idx < len(lms):
            out[key] = landmark_to_array(lms[idx])
    return out

async def handle_get_joint(ws: ServerConnection, data: Dict[str, Any]):
    manager = _manager_ref.get()
    if manager is None:
        await safe_send(ws, json.dumps({"type": "error", "message": "manager not ready"}))
        return

    person_id = int(data.get("person_id", 0))
    joint_name = data.get("joint")
    joint_idx = data.get("joint_idx")

    if joint_idx is None:
        if joint_name not in NAME_TO_IDX:
            await safe_send(ws, json.dumps({"type": "error", "message": f"unknown joint: {joint_name}"}))
            return
        joint_idx = NAME_TO_IDX[joint_name]

    lms = manager.get_latest_pose(person_id)
    if lms is None:
        await safe_send(ws, json.dumps({"type": "error", "message": f"person {person_id} not found"}))
        return

    if joint_idx < 0 or joint_idx >= len(lms):
        await safe_send(ws, json.dumps({"type": "error", "message": f"joint_idx out of range: {joint_idx}"}))
        return

    coords = landmark_to_array(lms[joint_idx])
    resp = {
        "type": "joint",
        "person_id": person_id,
        "joint": joint_name if joint_name is not None else joint_idx,
        "coords": coords,
        "ts": time.time(),
    }
    await safe_send(ws, json.dumps(resp, ensure_ascii=False))

async def handle_get_joints(ws: ServerConnection, data: Dict[str, Any]):
    manager = _manager_ref.get()
    if manager is None:
        await safe_send(ws, json.dumps({"type": "error", "message": "manager not ready"}))
        return

    person_id = int(data.get("person_id", 0))
    joints = data.get("joints", []) or DEFAULT_9

    # 多人查询：person_id < 0
    if person_id < 0:
        if not hasattr(manager, "get_all_latest_poses"):
            await safe_send(ws, json.dumps({"type": "error", "message": "manager missing get_all_latest_poses()"}))
            return
        poses = manager.get_all_latest_poses()  # Dict[int, List[landmark]]
        data_map = {}
        for pid, lms in poses.items():
            if lms is None:
                continue
            data_map[str(pid)] = collect_joints_from_landmarks(lms, joints)
        await safe_send(ws, json.dumps({
            "type": "joints",
            "person_id": -1,
            "data": data_map,
            "ts": time.time()
        }, ensure_ascii=False))
        return

    # 单人查询：person_id >= 0
    lms = manager.get_latest_pose(person_id)
    if lms is None:
        await safe_send(ws, json.dumps({"type": "error", "message": f"person {person_id} not found"}))
        return
    out = collect_joints_from_landmarks(lms, joints)
    await safe_send(ws, json.dumps({
        "type": "joints",
        "person_id": person_id,
        "data": out,
        "ts": time.time()
    }, ensure_ascii=False))

async def ws_handler(websocket: ServerConnection):
    CLIENTS.add(websocket)
    try:
        async for message in websocket:
            try:
                data = json.loads(message)
            except json.JSONDecodeError:
                await safe_send(websocket, json.dumps({"type": "error", "message": "invalid json"}, ensure_ascii=False))
                continue

            msg_type = data.get("type")
            try:
                if msg_type == "get_joint":
                    await handle_get_joint(websocket, data)
                elif msg_type == "get_joints":
                    await handle_get_joints(websocket, data)
                elif msg_type == "start_video":
                    # 可选：支持指定分辨率缩放/帧率等，当前直接订阅
                    SUBSCRIBERS.add(websocket)
                    await safe_send(websocket, json.dumps({"type": "ack", "op": "start_video"}))
                elif msg_type == "stop_video":
                    SUBSCRIBERS.discard(websocket)
                    await safe_send(websocket, json.dumps({"type": "ack", "op": "stop_video"}))                    
                else:
                    await safe_send(websocket, json.dumps({"type": "error", "message": f"unknown type: {msg_type}"}, ensure_ascii=False))
            except (websockets.exceptions.ConnectionClosed, asyncio.CancelledError):
                # 对端已断开，提前结束循环
                break
            except Exception as e:
                # 防止业务异常导致连接泄漏
                await safe_send(websocket, json.dumps({"type": "error", "message": f"server error: {e}"}, ensure_ascii=False))
    finally:
        CLIENTS.discard(websocket)
        SUBSCRIBERS.discard(websocket)

def get_subscribers():
    return SUBSCRIBERS

async def start_ws_server(host: Optional[str] = None, port: Optional[int] = None):
    host = host or HOST
    port = port or PORT
    server = await websockets.serve(ws_handler, host, port, ping_interval=30, ping_timeout=45, close_timeout=10)
    return server

__all__ = [
    "start_ws_server",
    "broadcast_json",
    "get_clients",
    "attach_manager",
    "send_video_frame_to_subscribers",
    "get_subscribers",
]